# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import logging
import subprocess
import numpy as np
import paddle
from collections import OrderedDict
import paddle.fluid as fluid
from paddle.fluid import core
from paddle.fluid.log_helper import get_logger

from google.protobuf import text_format
from paddle.fluid import debugger
from paddle.fluid.framework import Program
from paddle.fluid.proto import framework_pb2

__all__ = [
    "load_program",
    "save_program",
    "program_type_trans",
    "check_saved_vars_try_dump",
    "parse_program",
    "check_pruned_program_vars",
    "graphviz",
]

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
formatter = logging.Formatter(fmt='%(asctime)s - %(levelname)s - %(message)s')
ch = logging.StreamHandler()
ch.setFormatter(formatter)
logger.addHandler(ch)

persistable_vars_out_fn = "vars_persistable.log"
all_vars_out_fn = "vars_all.log"
ops_out_fn = "ops.log"

feed_fetch_type_list = [
    core.VarDesc.VarType.FEED_MINIBATCH,
    core.VarDesc.VarType.FETCH_LIST,
]
not_expected_op_types = ["lookup_table"]


def load_program(model_filename, is_text=False):
    if is_text:
        return load_program_text(model_filename)
    return load_program_binary(model_filename)


def load_program_binary(model_filename):
    """load program from binary string file"""
    with open(model_filename, "rb") as f:
        program_desc_str = f.read()
    return Program.parse_from_string(program_desc_str)


def load_program_text(model_filename):
    """load program from human-readable text file"""
    with open(model_filename, "r") as f:
        program_desc_text = f.read()

    prog_desc = framework_pb2.ProgramDesc()
    text_format.Merge(program_desc_text, prog_desc)
    return Program.parse_from_string(prog_desc.SerializeToString())


def save_program(program, model_filename='__model__', is_text=False):
    if is_text:
        with open(model_filename, "w") as f:
            f.write(str(program))
    else:
        with open(model_filename, "wb") as f:
            f.write(program.desc.serialize_to_string())


def check_pruned_program_vars(train_prog, pruned_prog):
    is_match = True

    pruned_vars = [
        (v.name, v)
        for v in pruned_prog.list_vars()
        if fluid.io.is_persistable(v)
    ]
    pruned_vars = OrderedDict(pruned_vars)
    pruned_vars_name = [name for name in pruned_vars]
    logger.info(
        "persistable vars in pruned program: {}".format(pruned_vars_name)
    )

    for var_name in pruned_vars:
        var = pruned_vars[var_name]
        # feed and fetch op is added in pruned program when pruning, not need to be found in train program
        if var.type in feed_fetch_type_list:
            break
        try:
            train_prog_var = train_prog.global_block().var(var_name)
        except ValueError as e:
            logger.error(
                "not find variable '%s' in train program. please check pruning."
                % var_name
            )
            logger.error(e)
            continue
        if (
            var.shape != train_prog_var.shape
            or var.dtype != train_prog_var.dtype
        ):
            logger.error(
                "variable: {} not match. in pruned program shape: {} dtype:{}, in train program shape: {} dtype: {}".format(
                    var_name,
                    var.shape,
                    var.dtype,
                    train_prog_var.shape,
                    train_prog_var.dtype,
                )
            )
            is_match = False
    return is_match


def graphviz(block, output_dir="", filename='debug'):
    dot_path = os.path.join(output_dir, filename + '.dot')
    pdf_path = os.path.join(output_dir, filename + '.pdf')
    debugger.draw_block_graphviz(block, path=dot_path)
    cmd = ["dot", "-Tpdf", dot_path, "-o", pdf_path]
    p = subprocess.Popen(
        cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    p.wait()


def program_type_trans(prog_dir, prog_fn, is_text):
    prog = load_program(os.path.join(prog_dir, prog_fn), is_text)
    prog_out_fn = prog_fn + ".bin" if is_text else prog_fn + ".pbtxt"
    save_program(prog, os.path.join(prog_dir, prog_out_fn), 1 - is_text)
    return prog_out_fn


def append_save_op(block, var, path):
    block.append_op(
        type='save', inputs={'X': [var]}, outputs={}, attrs={'file_path': path}
    )


def append_load_op(block, var, path):
    block.append_op(
        type='load',
        inputs={},
        outputs={'Out': [var]},
        attrs={'file_path': path},
    )


def save_var(np_array, var_name, shape_list, dtype, save_path):
    program = fluid.Program()
    place = fluid.CPUPlace()
    exe = fluid.Executor(place)
    shape = list(shape_list)
    with fluid.program_guard(program):
        d0_data = paddle.static.data(var_name, shape=shape, dtype=dtype)
        append_save_op(program.global_block(), d0_data, save_path)
        exe.run(feed={var_name: np_array}, fetch_list=[])


def load_var(var_name, shape_list, dtype, save_path):
    program = fluid.Program()
    place = fluid.CPUPlace()
    exe = fluid.Executor(place)
    with fluid.program_guard(program):
        d0_data = paddle.static.data(var_name, shape=shape_list, dtype=dtype)
        append_load_op(program.global_block(), d0_data, save_path)
        outs = exe.run(feed={}, fetch_list=[d0_data])
        return outs


def reader(batch_size, fn, dim):
    data = []
    if isinstance(dim, list) or isinstance(dim, tuple):
        shape = list(dim)
        _temp = 1
        for x in dim:
            _temp = _temp * x
        dim = _temp
    else:
        shape = [dim]

    shape = [batch_size] + shape
    dim = dim * batch_size

    for line in open(fn, 'r'):
        fields = line.strip().split(' ')
        fields = [float(d) for d in fields]
        while len(fields) >= dim:
            tmp = fields[:dim]
            fields = fields[dim:]
            data.append(np.array(tmp).reshape(shape))
    return data


def feed_gen(batch_size, feeded_vars_dims, feeded_vars_filelist):
    batch_feed = []
    for i, fn in enumerate(feeded_vars_filelist):
        batch_feed.append(reader(batch_size, fn, feeded_vars_dims[i]))
    return batch_feed


def try_load_model_vars(
    dump_dir,
    dump_prog_fn,
    is_text_dump_program,
    batch_size,
    feed_config,
    fetch_config,
    save_filename,
    saved_params,
):
    place = fluid.CPUPlace()
    exe = fluid.Executor(place)
    scope = fluid.core.Scope()
    with fluid.scope_guard(scope):
        if is_text_dump_program:
            dump_prog_fn = program_type_trans(
                dump_dir, dump_prog_fn, is_text_dump_program
            )
        (
            inference_program,
            feed_target_names,
            fetch_targets,
        ) = fluid.io.load_inference_model(
            dump_dir,
            exe,
            model_filename=dump_prog_fn,
            params_filename=save_filename,
        )

        # check program vars and saved vars shape
        orig_para_shape = {
            each_var.name: tuple(each_var.desc.shape())
            for each_var in saved_params
        }
        for each_var in saved_params:
            var_temp = fluid.global_scope().find_var(each_var.name)
            assert var_temp is not None, "can't not find var: " + each_var.name
            new_shape = (np.array(var_temp.get_tensor())).shape
            assert each_var.name in orig_para_shape, (
                each_var.name + "MUST in var list"
            )
            orig_shape = orig_para_shape.get(each_var.name)
            if new_shape != orig_shape:
                raise RuntimeError(
                    "Shape not matching: the Program requires a parameter with a shape of ({}), "
                    "while the loaded parameter (namely [ {} ]) has a shape of  ({}).".format(
                        orig_shape, each_var.name, new_shape
                    )
                )

        # check feed/fetch vars in program and config
        fetch_targets_names = [v.name for v in fetch_targets]
        if not feed_target_names:
            logger.warning("no feed targets in program.")
        if not fetch_targets_names:
            logger.warning("no fetch targets in program.")
        fetch_list = fetch_targets
        feed_name_list = feed_target_names
        if (
            feed_config.feeded_vars_names is not None
            and feed_target_names != feed_config.feeded_vars_names
        ):
            logger.warning(
                "feed vars in program and config are diff: feed in program: {}. feed in config {}.".format(
                    feed_target_names, feed_config.feeded_vars_names
                )
            )
            feed_name_list = feed_config.feeded_vars_names
            # remove feed op in inference_program. new feed op will be added in exe.run
            global_block = inference_program.global_block()
            need_to_remove_op_index = []
            for i, op in enumerate(global_block.ops):
                op.desc.set_is_target(False)
                if op.type == "feed":  # only remove feed op here
                    need_to_remove_op_index.append(i)
            for index in need_to_remove_op_index[::-1]:
                global_block._remove_op(index)
        if (
            fetch_config.fetch_vars_names is not None
            and fetch_targets_names != fetch_config.fetch_vars_names
        ):
            logger.warning(
                "fetch vars in program and config are diff: fetch in program: {}. fetch in config {}.".format(
                    fetch_targets_names, fetch_config.fetch_vars_names
                )
            )
            fetch_list = [
                inference_program.global_block().var(i)
                for i in fetch_config.fetch_vars_names
            ]
            # remove fetch op in inference_program. new fetch op will be added in exe.run
            global_block = inference_program.global_block()
            need_to_remove_op_index = []
            for i, op in enumerate(global_block.ops):
                op.desc.set_is_target(False)
                if op.type == "fetch":  # only remove fetch op here
                    need_to_remove_op_index.append(i)
            for index in need_to_remove_op_index[::-1]:
                global_block._remove_op(index)

        # if fetch_list have lod tensor
        return_numpy = all([v.lod_level == 0 for v in fetch_list])

        # try dump fetch_targets
        feed_tensors = []
        assert (
            len(feed_config.feeded_vars_names)
            == len(feed_config.feeded_vars_dims)
            == len(feed_config.feeded_vars_types)
        )
        # check program vars and feed tensor shape in config
        for i in range(len(feed_config.feeded_vars_names)):
            var = inference_program.global_block().var(
                feed_config.feeded_vars_names[i]
            )
            if not isinstance(feed_config.feeded_vars_dims[i], (list, tuple)):
                tensor_shape = (feed_config.feeded_vars_dims[i],)
            else:
                tensor_shape = tuple(feed_config.feeded_vars_dims[i])
            feed_config.feeded_vars_dims[i] = tensor_shape
            var_shape = var.shape[1:]
            if tensor_shape != var_shape:
                raise RuntimeError(
                    "feed variable '{}' shape not match. infer program  shape: {}. feed tensor shape: {}".format(
                        feed_config.feeded_vars_names[i],
                        var_shape,
                        tensor_shape,
                    )
                )

        if not feed_config.feeded_vars_filelist:
            logger.info("generate random feed vars.")
            for i in range(len(feed_config.feeded_vars_names)):
                var = inference_program.global_block().var(
                    feed_config.feeded_vars_names[i]
                )
                # create fake feed tensor. if lod_level > 1, should create_lod_tensor()
                if var.lod_level == 0:
                    feed_tensors.append(
                        np.array(
                            np.random.random(
                                tuple(
                                    [batch_size]
                                    + list(feed_config.feeded_vars_dims[i])
                                )
                            ),
                            dtype=feed_config.feeded_vars_types[i],
                        )
                    )
                elif var.lod_level == 1:
                    t = np.array(
                        np.random.random(
                            tuple(
                                [batch_size]
                                + list(feed_config.feeded_vars_dims[i])
                            )
                        ),
                        dtype=feed_config.feeded_vars_types[i],
                    )
                    feed_tensors.append(
                        fluid.create_lod_tensor(t, [[1] * batch_size], place)
                    )
                else:
                    raise RuntimeError(
                        "vars with lod_level >= 2 is not supported now in this infer program check tool."
                    )
            results = exe.run(
                inference_program,
                feed={
                    name: feed_tensors[i]
                    for i, name in enumerate(feed_name_list)
                },
                fetch_list=fetch_list,
                return_numpy=return_numpy,
            )
        else:
            logger.info(
                "load feed vars from files: {}.".format(
                    feed_config.feeded_vars_filelist
                )
            )
            feed_vars = [
                inference_program.global_block().var(
                    feed_config.feeded_vars_names[i]
                )
                for i in range(len(feed_config.feeded_vars_names))
            ]
            feeder = fluid.DataFeeder(feed_list=feed_vars, place=place)
            batch_feed = feed_gen(
                batch_size,
                feed_config.feeded_vars_dims,
                feed_config.feeded_vars_filelist,
            )
            slots = [batch_feed]
            results = exe.run(
                inference_program,
                feed=feeder.feed(slots),
                fetch_list=fetch_list,
                return_numpy=return_numpy,
            )
        for i, v in enumerate(fetch_list):
            logger.info("fetch_targets name: %s" % v.name)
            logger.info("fetch_targets: {}".format(results[i]))
        return results


def check_not_expected_ops(prog):
    op_types_set = set()
    for op in prog.global_block().ops:
        if op.type in not_expected_op_types and op.type not in op_types_set:
            logger.warning(
                "find op type '{}' in program, please check if your program is pruned correctly !".format(
                    op.type
                )
            )
            op_types_set.add(op.type)


def check_saved_vars_try_dump(
    dump_dir,
    dump_prog_fn,
    is_text_dump_program,
    feed_config,
    fetch_config,
    batch_size=1,
    save_filename=None,
):
    dump_prog = load_program(
        os.path.join(dump_dir, dump_prog_fn), is_text_dump_program
    )
    saved_params = [
        v for v in dump_prog.list_vars() if fluid.io.is_persistable(v)
    ]
    logger.info(
        "persistable vars in dump program: {}".format(
            [v.name for v in saved_params]
        )
    )

    check_not_expected_ops(dump_prog)

    return try_load_model_vars(
        dump_dir,
        dump_prog_fn,
        is_text_dump_program,
        batch_size,
        feed_config,
        fetch_config,
        save_filename,
        saved_params,
    )


def parse_program(program, output_dir):
    # persistable vars
    output = {}
    persistable_vars = [
        v for v in program.list_vars() if fluid.io.is_persistable(v)
    ]
    output["persistable_vars"] = [
        {
            'name': str(v.name),
            'shape': str(v.shape),
            'lod_level': int(v.lod_level),
            'dtype': str(v.dtype),
            'type': str(v.type),
        }
        for v in persistable_vars
    ]
    with open(os.path.join(output_dir, persistable_vars_out_fn), 'w') as f:
        f.write("persistable vars:\n")
        for var in output["persistable_vars"]:
            f.write(str(var))
            f.write("\n")

    # all vars
    all_vars = [v for v in program.list_vars()]
    output["all_vars"] = [
        {
            'name': str(v.name),
            'shape': str(v.shape),
            'lod_level': int(v.lod_level),
            'dtype': str(v.dtype),
        }
        if v.type not in feed_fetch_type_list
        else {'name': str(v.name), 'type': str(v.type)}
        for v in all_vars
    ]
    with open(os.path.join(output_dir, all_vars_out_fn), 'w') as f:
        f.write("all vars:\n")
        for var in output["all_vars"]:
            f.write(str(var))
            f.write("\n")

    # ops
    ops = program.global_block().ops
    output["ops"] = [
        {
            'type': op.type,
            'input_arg_names': str(op.input_arg_names),
            'output_arg_names': str(op.output_arg_names),
        }
        for op in ops
    ]
    with open(os.path.join(output_dir, ops_out_fn), 'w') as f:
        f.write("ops:\n")
        for op in output["ops"]:
            f.write(str(op))
            f.write("\n")
